#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

std::vector<cv::Vec4i> _probabilistic_hough_lines(const cv::Mat& binary, float rho, float theta, int threshold, int lineLength, int lineGap) {
    CV_Assert(binary.type() == CV_8UC1);
    
    int width = binary.cols;
    int height = binary.rows;
    
    int numangle = cvRound(CV_PI / theta);
    int max_rho = (int)std::sqrt(width*width + height*height) + 1;
    int min_rho = -max_rho;
    int numrho = cvRound((max_rho - min_rho + 1) / rho);
    
    cv::Mat accum = cv::Mat::zeros(numangle, numrho, CV_32SC1);
    std::vector<float> cosTable(numangle), sinTable(numangle);
    for (int i = 0; i < numangle; i++) {
        cosTable[i] = (float)(std::cos((double)i*theta));
        sinTable[i] = (float)(std::sin((double)i*theta));
    }
    
    std::vector<cv::Point> nonZeroLocations;
    cv::Mat nonZeroMat = (binary != 0);
    cv::findNonZero(nonZeroMat, nonZeroLocations);
    cv::Mat markers = cv::Mat::zeros(height, width, CV_8UC1);
    for (int i = 0; i < nonZeroLocations.size(); i++) {
        cv::Point pt = nonZeroLocations[i];
        markers.ptr<uchar>(pt.y)[pt.x] = (uchar)1;
    }
    
    int count = nonZeroLocations.size();
    cv::RNG rng((uint64)-1);
    std::vector<cv::Vec4i> lines;
    while (count > 0) {
        count--;
        
        int idx = rng.uniform(0, count+1);
        int max_val = threshold-1;
        int max_n = 0;
        
        cv::Point point = nonZeroLocations[idx];
        nonZeroLocations[idx] = nonZeroLocations[count];
        
        if (markers.ptr<uchar>(point.y)[point.x] == 0) {
            continue;
        }
        
        for (int n = 0; n < numangle; n++) {
            int r = cvRound((point.x * cosTable[n] + point.y * sinTable[n]) / rho);
            r += (numrho - 1) / 2;
            accum.ptr<int>(n)[r]++;
            if (accum.ptr<int>(n)[r] > max_val) {
                max_val = accum.ptr<int>(n)[r];
                max_n = n;
            }
        }
        
        if (max_val < threshold) {
            continue;
        }
        
        float a = -sinTable[max_n];
        float b = cosTable[max_n];
        int xflag, dx0, dy0;
        int x0 = point.x, y0 = point.y;
        if (std::fabs(a) > std::fabs(b)) {
            xflag = 1;
            dx0 = -1;
            dy0 = cvRound(b * (1 << 16) / std::fabs(a));
            y0 = (y0 << 16) + (1 << 15);
        } else {
            xflag = 0;
            dy0 = b > 0 ? 1 : -1;
            dx0 = cvRound(a * (1 << 16) / std::fabs(b));
            x0 = (x0 << 16) + (1 << 15);
        }
        
        cv::Point line_end[2];
        for (int  k = 0; k < 2; k++) {
            int gap = 0, x = x0, y = y0, dx = dx0, dy = dy0;
            if (k > 0) {
                dx = -dx;
                dy = -dy;
            }
            
            for (;; x += dx, y += dy) {
                int i1, j1;
                if (xflag) {
                    j1 = x;
                    i1 = y >> 16;
                } else {
                    j1 = x >> 16;
                    i1 = y;
                }
                
                if (j1 < 0 || j1 >= width || i1 < 0 || i1 >= height) {
                    break;
                }
                
                if (markers.ptr<uchar>(i1)[j1] != 0) {
                    gap = 0;
                    line_end[k].y = i1;
                    line_end[k].x = j1;
                } else {
                    gap++;
                    if (gap > lineGap) {
                        break;
                    }
                }
            }
        }
        
        bool good_line = (std::abs(line_end[1].x - line_end[0].x) >= lineLength || std::abs(line_end[1].y - line_end[0].y) >= lineLength);
        
        for (int k = 0; k < 2; k++) {
            int x = x0, y = y0, dx = dx0, dy = dy0;
            if (k > 0) {
                dx = -dx;
                dy = -dy;
            }
            for (;; x+=dx, y+=dy) {
                int i1, j1;
                if (xflag) {
                    j1 = x;
                    i1 = y >> 16;
                } else {
                    j1 = x >> 16;
                    i1 = y;
                }
                if (j1 < 0 || j1 >= width || i1 < 0 || i1 >= height) {
                    continue;
                }
                
                if (markers.ptr<uchar>(i1)[j1] != 0) {
                    if (good_line) {
                        for (int n = 0; n < numangle; n++) {
                            int r = cvRound((j1 * cosTable[n] + i1 * sinTable[n]) / rho);
                            r += (numrho - 1) / 2;
                            accum.ptr<int>(n)[r]--;
                        }
                    }
                    markers.ptr<uchar>(i1)[j1] = 0;
                }
                
                if (i1 == line_end[k].y && j1 == line_end[k].x) {
                    break;
                }
            }
        }
        
        if (good_line) {
            cv::Vec4i line(line_end[0].x, line_end[0].y, line_end[1].x, line_end[1].y);
            lines.push_back(line);
        }
    }
    return lines;
}

bool _lines_are_same(const std::vector<cv::Vec4i>& lines1, const std::vector<cv::Vec4i>& lines2) {
    if (lines1.size() != lines2.size()) {
        return false;
    }
    for (int i = 0; i < lines1.size(); i++) {
        if (lines1[i][0] != lines2[i][0] || lines1[i][1] != lines2[i][1]
            || lines1[i][2] != lines2[i][2] || lines1[i][3] != lines2[i][3]) {
            return false;
        }
    }
    return true;
}

int main(int argc, char **argv) {
    cv::Mat src = cv::imread("/home/xlll/Downloads/opencv/samples/data/building.jpg", cv::IMREAD_COLOR);
    if (src.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat gray;
    cv::Canny(src, gray, 50, 200, 3);
    cv::Mat binary1 = gray.clone();
    cv::Mat binary2 = gray.clone();
    
    std::vector<cv::Vec4i> lines;
    double rho = 1;
    double theta = CV_PI / 180;
    double threshold = 200;
    double minLineLength = 80;
    double maxLineGap = 10;
    cv::HoughLinesP(binary1, lines, rho, theta, threshold, minLineLength, maxLineGap);
    std::cout << "number of lines = " << lines.size() << std::endl;
    cv::Mat dst = src.clone();
    for (int i = 0; i < lines.size(); i++) {
        cv::line(dst, cv::Point(lines[i][0], lines[i][1]), cv::Point(lines[i][2], lines[i][3]), cv::Scalar(0, 0, 255), 3, cv::LINE_AA);
    }
    std::cout << "==========" << std::endl;
    
    std::vector<cv::Vec4i> lines2;
    lines2 = _probabilistic_hough_lines(binary2, rho, theta, threshold, minLineLength, maxLineGap);
    std::cout << "number of lines = " << lines2.size() << std::endl;
    
    std::cout << (_lines_are_same(lines, lines2) ? "Results are the same!" : "Results are not the same!") << std::endl;
    
    cv::imshow("src", src);
    cv::imshow("binary", binary1);
    cv::imshow("dst", dst);
    cv::waitKey(0);
    
    return EXIT_SUCCESS;
}
